import os

import numpy as np
import torch
import torch.nn as nn
# from dataset import RSCDataset, train_transform, val_transform
from PIL import Image
from torch.cuda.amp import autocast

from core import train_net
from core.builder import timm_builder
from core.config import cfg
from core.data import Data_Pipeline, train_transform, val_transform

Image.MAX_IMAGE_PIXELS = 1000000000000000

cfg_file_path = 'configs/mnist_res18.yaml'

# 参数设置
if cfg_file_path != '':
    cfg.merge_from_file(cfg_file_path)
cfg.freeze()

#os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.SYSTEM.GPUS_INDEX
device = torch.device("cuda")

# 准备数据集
train_imgs_dir = os.path.join(cfg.TRAIN.DATA_PATH, "mnist_data/train/")
val_imgs_dir = os.path.join(cfg.TRAIN.DATA_PATH, "mnist_data/test/")

train_labels_dir = os.path.join(cfg.TRAIN.DATA_PATH, "mnist_datatrain.txt")
val_labels_dir = os.path.join(cfg.TRAIN.DATA_PATH, "mnist_datatest.txt")

train_data = Data_Pipeline(train_imgs_dir, train_labels_dir, transform=train_transform)
valid_data = Data_Pipeline(val_imgs_dir, val_labels_dir, transform=val_transform)

model = timm_builder(cfg.TRAIN.MODEL_NAME,cfg.TRAIN.N_CLASS).cuda()
model= torch.nn.DataParallel(model)
# checkpoints=torch.load('outputs/efficientnet-b6-3729/ckpt/checkpoint-epoch20.pth')
# model.load_state_dict(checkpoints['state_dict'])

# 模型保存路径
save_ckpt_dir = os.path.join(cfg.TRAIN.SAVE_CKPT_PATH, cfg.TRAIN.MODEL_NAME)
save_log_dir = os.path.join(cfg.TRAIN.LOG_PATH , cfg.TRAIN.MODEL_NAME)

if not os.path.exists(save_ckpt_dir):
    os.makedirs(save_ckpt_dir)
if not os.path.exists(save_log_dir):
    os.makedirs(save_log_dir)

# 训练
best_model, model = train_net(cfg, model, train_data, valid_data)

